import pysam
import numpy as np
from deepsignal3.utils.process_utils import complement_seq
from deepsignal3.utils.ref_reader import get_contig2len,get_contig2len_n_seq

from deepsignal3.utils import bam_reader

import gzip
import time
from multiprocessing import Manager
import multiprocessing as mp

def process_deepmod2(features_file, feature_Q, existing_keys=None, time_wait=1, reads_per_chunk=2, qsize_limit=20):
    if features_file.endswith(".gz"):
        infile = gzip.open(features_file, 'rt')
    else:
        infile = open(features_file, 'r')
    pred_deepmod2 = []
    fea_list = []
    pre_read_id = ''
    chromosome = ''
    read_id = ''
    num=0
    # encodings = ['utf-8', 'gbk', 'latin-1']
    for line in infile:
        # if features_file.endswith(".gz"):
        #    words = line.decode('gbk').strip().split("\t")
        # else:
        #    words=line.strip().split("\t")
        words = line.strip().split("\t")
        read_id = words[0]
        if pre_read_id == '':
            pre_read_id = read_id
        elif pre_read_id != read_id:
            fea_list.append((pre_read_id, chromosome, pred_deepmod2))
            pre_read_id = read_id
            pred_deepmod2 = []
            if len(fea_list) >= reads_per_chunk:
                while feature_Q.qsize() >= qsize_limit:
                    # print('Pausing tsv input due to INPUT queue size limit. Signal_qsize=%d' %(feature_Q.qsize()), flush=True)
                    time.sleep(time_wait)
                feature_Q.put(fea_list)
                fea_list = []

        chromosome = words[1]
        pos=words[2]
        key='||'.join([pre_read_id,chromosome,words[2]])
        if existing_keys is not None:
            if key in existing_keys:
                num+=1
                continue
        if pos=='NA':
            continue
        pred_deepmod2.append(int(pos))
    if pre_read_id != read_id:
        fea_list.append((pre_read_id, chromosome, pred_deepmod2))
    if len(fea_list) > 0:
        feature_Q.put(fea_list)
    infile.close()
    print('Number of intersection points: {}'.format(num))

def process_bam(bam_index,feature_Q, time_wait=1):
    case_num=0
    ref_unmatch_num=0
    filter=0
    while True:
        while feature_Q.empty():
            time.sleep(time_wait)
        read_data = feature_Q.get()
        if read_data == "kill":
            feature_Q.put("kill")
            break
        for (read_name, chromosome, pred_deepmod2) in read_data:
            try:
                flag=0
                for bam_read in bam_index.get_alignments(read_name):
                    if bam_read.is_supplementary or bam_read.is_secondary:
                        continue
                    flag=1
                    reference_name = bam_read.reference_name
                    if reference_name!=chromosome:
                        print('Comparing chromosomes is different of readid {}'.format(read_name))                    
                        continue
                    strand_code = 0 if bam_read.is_reverse else 1
                    strand = "-" if strand_code == 0 else "+"
                    ref_start = bam_read.reference_start
                    ref_end = bam_read.reference_end

                    seq = bam_read.get_forward_sequence()
                    if seq is None:
                        print('not have origin seq of readid {}'.format(read_name))
                        continue
                    if reference_name is None:
                        print('not have reference name of readid {}'.format(read_name))
                        continue
                    if bam_read.infer_query_length() != len(seq):
                        print(
                            'infer read length is not same as seq length of readid {}'.format(read_name))
                        continue
                    if not bam_read.has_tag('MD'):
                        print('not have MD of readid {}'.format(read_name)) 
                        continue
                    if bam_read.is_reverse:
                        ref_seq = complement_seq(bam_read.get_reference_sequence())
                    else:
                        ref_seq = bam_read.get_reference_sequence()
                    for ref_pos in pred_deepmod2:
                        if strand_code == 0:
                            pos=ref_end-ref_pos-1
                            if ref_seq[pos:pos+2]!='CG':
                                if ref_seq[pos:pos+2].upper()=='CG':
                                    case_num+=1
                                else:
                                    ref_unmatch_num+=1
                        else:
                            pos=ref_pos-ref_start
                            if ref_seq[pos:pos+2]!='CG':
                                if ref_seq[pos:pos+2].upper()=='CG':
                                    case_num+=1
                                else:
                                    ref_unmatch_num+=1
                if flag==0:
                    filter+=len(pred_deepmod2)
            except:
                print('Error in readid {}'.format(read_name))
    print('Number of cases: {}'.format(case_num))
    print('Number of reference unmatch: {}'.format(ref_unmatch_num))


if __name__ == '__main__':
    # bamfile = pysam.AlignmentFile('/home/xiaoyifu/data/HG002/R9.4/part2/bam/dorado.5mc_5hmc.sup.sort.bam', "rb", check_sq=False)
    # bam_index=pysam.IndexedReads(bamfile)
    # bam_index.build()
    manager = mp.Manager()
    feature_Q = manager.Queue()
    bam_index = bam_reader.ReadIndexedBam('/home/xiaoyifu/data/HG002/R9.4/part2/bam/dorado.5mc_5hmc.sup.sort.bam')
    pd = mp.Process(target=process_deepmod2, args=('/homeb/xiaoyf/HG002/part2/deepmod2_t5_uniq.keys.tsv', feature_Q), name="fea_reader")
    pd.daemon = True
    pd.start()
    pb = mp.Process(target=process_bam, args=(bam_index,  feature_Q),
                        name="pb_reader")

    pb.daemon = True
    pb.start()
    pd.join()
    feature_Q.put("kill")
    pb.join()